
import time
import torch
import numpy
import random

from datetime import datetime
from babel.dates import format_date
from torch.nn.functional import one_hot


class PseudoDateDataset(torch.utils.data.Dataset):

    def __init__(self, size: int, pad_token: str = '<pad>', unk_token: str = '<unk>') -> None:
        """
        :param size		: 数据集大小.
        :param pad_token: Token of padding words.
        :param unk_token: Token of stopping words.
        """
        human_vocabulary = set()
        machine_vocabulary = set()
        self.dataset = []
        for _ in range(size):
            human_readable_date, machine_readable_date = self.generate_pseudo_date()
            self.dataset.append((human_readable_date, machine_readable_date))
            human_vocabulary.update(tuple(human_readable_date))
            machine_vocabulary.update(tuple(machine_readable_date))
        self.human_vocabulary = dict(
            zip(sorted(human_vocabulary) + [unk_token, pad_token], list(range(len(human_vocabulary) + 2))))
        self.machine_vocabulary = {token: index for index, token in enumerate(sorted(machine_vocabulary) + [pad_token])}
        self.inverse_machine_vocabulary = {index: token for token, index in self.machine_vocabulary.items()}
        self.human_unk = self.human_vocabulary.get(unk_token)  # 在人为词汇中获得'<unk>'标记的id.
        self.machine_unk = self.machine_vocabulary.get(unk_token)  # 在机器词汇中获得'<unk>'标记的id.
        self.machine_pad = self.machine_vocabulary.get(pad_token)  # 在机器词汇中获得'<pad>'标记的id
        self.pad_token = pad_token
        self.unk_token = unk_token

    def generate_pseudo_date(self, builtin_ratio: float = .8, builtin_dist: list = [.2, .2, .2, .4],
                             locale: str = 'en_US') -> (str, str):
        """
        生成一对人类可读的日期字符串和机器可读的日期字符串.
        :param builtin_ratio			: 尽可能从' builtin_datetime_formats '而不是' custom_datetime_formats '中选择日期时间格式
        :param built_dist				: 在' builtin_datetime_formats '(短，中，长，全)中选择不同日期时间格式的可能性分布
        :param locale					: 日期时间的语言，例如:'en_US'， 'zh_CN'，默认的'en_US'表示美式英语。

        :return human_readable_date		: 例如 'Wednesday, October 26, 2107'.
        :return machine_readable_date	: 时间日期格式 'yyyy-MM-dd'
        """

        # 随机生成人类可读的日期格式
        def _generate_datetime_format(builtin_ratio: float, builtin_dist: list) -> str:
            builtin_datetime_formats = [
                'short',  # d/MM/YY
                'medium',  # MMM d, YYYY
                'long',  # MMMM, d, YYYY
                'full',  # EEEE, MMMM d, YYYY
            ]
            custom_datetime_formats = [  # 一些常见的人类可读日期时间格式可以按惯例枚举.
                '{month}/{day}/{year}'.format,
                '{year}/{month}/{day}'.format,
                '{month} {day}, {year}'.format,
                '{year}.{month}.{day}'.format,
                '{day} {month}, {year}'.format,
                '{year}, {day} {month}'.format,
            ]
            year_formats = ['Y', 'YY', 'YYY', 'YYYY']  # 2020 20 2020 2020
            month_formats = ['M', 'MM', 'MMM', 'MMMM']  # 2 02 Feb February
            day_formats = ['d', 'dd']  # 2 02
            weekday_formats = ['E', 'E', 'EEE', 'EEEE']  # Sat Sat Sat Saturday
            if random.random() < builtin_ratio:
                return numpy.random.choice(a=builtin_datetime_formats, p=builtin_dist)
            return random.choice(custom_datetime_formats)(year=random.choice(year_formats),
                                                          month=random.choice(month_formats),
                                                          day=random.choice(day_formats))

        random_timestamp = int(
            random.random() * time.time()) + 86400  # 随机生成一个时间戳:注意' datetime.datetime.fromtimestamp '接收的时间戳大于86399。
        random_date = datetime.fromtimestamp(random_timestamp).date()  # 随机生成一个时间
        datetime_format = _generate_datetime_format(builtin_ratio=builtin_ratio, builtin_dist=builtin_dist)
        human_readable_date = format_date(random_date, format=datetime_format, locale=locale).lower().replace(',', '')
        machine_readable_date = random_date.isoformat()
        return human_readable_date, machine_readable_date

    def __getitem__(self, idx: int) -> (torch.LongTensor, torch.LongTensor, dict):
        human_readable_date, machine_readable_date = self.dataset[idx]  # 在dataset中获取元组
        human_readable_vector = list(map(lambda x: self.human_vocabulary.get(x, self.human_unk), human_readable_date))
        machine_readable_vector = [self.machine_pad] + list(
            map(lambda x: self.machine_vocabulary.get(x, self.machine_unk), machine_readable_date)) + [self.machine_pad]
        human_readable_tensor = one_hot(torch.tensor(human_readable_vector), num_classes=len(self.human_vocabulary))
        machine_readable_tensor = one_hot(torch.tensor(machine_readable_vector),
                                          num_classes=len(self.machine_vocabulary))
        raw = {'human_readable_date': human_readable_date, 'machine_readable_date': machine_readable_date}
        return human_readable_tensor, machine_readable_tensor, raw

    def __len__(self) -> int:
        return len(self.dataset)


if __name__ == '__main__':
    from utils import collate_fn

    dataset = PseudoDateDataset(size=1000)
    padding_tensor = torch.FloatTensor(numpy.zeros(len(dataset.human_vocabulary)))
    padding_tensor[dataset.human_vocabulary.get(dataset.pad_token)] = 1
    dataloader = torch.utils.data.DataLoader(
        dataset=dataset,
        batch_size=4,
        shuffle=True,
        num_workers=0,
        collate_fn=collate_fn(padding_tensor),
    )
    for step, (batch_x, batch_y, batch_raw) in enumerate(dataloader):
        pass
